//这个函数使用对称高斯赛德尔迭代，来解方程组Ax=b，输入是A与b，输出是x
//具体是怎么迭代的，可以查看这个目录下的文档。

#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<cusparse.h>
#include"Matrix_segmentation.h"


//我们要解下面的这么一个方程组
/*
4x-y+z=7
4x−8y+z=−21
−2x+y+5z=15
*/



int main(){
    //首先我们先得到矩阵A的csr格式
    //矩阵大小为3
    int A_size = 3;
    //矩阵非零元素为9
    int nz = 9;

    //需要三个向量，一个是row、col、val
    int* A_row_host = 0;
    int* A_col_host = 0;
    float* A_val_host = 0;
    // printf("申请A空间\n");

    cudaMallocHost((void **)&A_col_host, nz * sizeof(int));
    cudaMallocHost((void **)&A_val_host, nz * sizeof(float));
    cudaMallocHost((void **)&A_row_host, (A_size + 1) * sizeof(int));

    //根据上面的方程组初始化
    // printf("根据上面的方程组初始化\n");

    A_row_host[0] = 0;
    A_row_host[1] = 3;
    A_row_host[2] = 6;
    A_row_host[3] = 9;

    A_col_host[0] = 0;
    A_col_host[1] = 1;
    A_col_host[2] = 2;
    A_col_host[3] = 0;
    A_col_host[4] = 1;
    A_col_host[5] = 2;
    A_col_host[6] = 0;
    A_col_host[7] = 1;
    A_col_host[8] = 2;

    A_val_host[0] = 4;
    A_val_host[1] = -1;
    A_val_host[2] = 1;
    A_val_host[3] = 4;
    A_val_host[4] = -8;
    A_val_host[5] = 1;
    A_val_host[6] = -2;
    A_val_host[7] = 1;
    A_val_host[8] = 5;

    //然后初始化x与b，我们让x的初值从0开始
    // printf("然后初始化x与b，我们让x的初值从0开始\n");
    float* x_val_host = 0;

    cudaMallocHost((void **)&x_val_host, A_size * sizeof(float));

    int i;
    for(i = 0; i < A_size; i++){
        x_val_host[i] = 0;
    }

    float* b_val_host = 0;

    cudaMallocHost((void **)&b_val_host, A_size * sizeof(float));

    b_val_host[0] = 7;
    b_val_host[1] = -21;
    b_val_host[2] = 15;

    //获取A的上三角部分
    //创造矩阵U
    // printf("创造矩阵U\n");

    //获取A的下三角部分
    int L_nz = 0;
    int* L_row_host = 0;
    float* L_val_host = 0;
    int* L_col_host = 0;

    cudaMallocHost((void **)&L_row_host, (A_size + 1) * sizeof(int));
    cudaMallocHost((void **)&L_val_host, nz * sizeof(float));
    cudaMallocHost((void **)&L_col_host, nz * sizeof(int));

    martix_csr_get_tril(A_val_host, A_row_host, A_col_host, nz, A_size,
        L_val_host, L_row_host, L_col_host, &L_nz);

    //获取A的D+U部分
    int DU_nz = 0;
    int* DU_row_host = 0;
    float* DU_val_host = 0;
    int* DU_col_host = 0;

    cudaMallocHost((void **)&DU_row_host, (A_size + 1) * sizeof(int));
    cudaMallocHost((void **)&DU_val_host, nz * sizeof(float));
    cudaMallocHost((void **)&DU_col_host, nz * sizeof(int));

    martix_csr_get_diag_triu(A_val_host, A_row_host, A_col_host, nz, A_size,
        DU_val_host, DU_row_host, DU_col_host, &DU_nz);

    //这里进行高斯赛德尔迭代
    //首先使用cusparseDcsrsv_solve
    //我们需要求解(D + U)x_{k+1} = b - Lx_{k}方程，我们需要分别把(D + U)与b - Lx_{k}搞定
    //先只进行一次迭代
    //https://blog.csdn.net/mifangdebaise/article/details/42811651
    
    //创建L矩阵的GPU空间
    int* L_row_device = 0;
    float* L_val_device = 0;
    int* L_col_device = 0;

    cudaMalloc((void **)&L_row_device, (A_size + 1) * sizeof(int));
    cudaMalloc((void **)&L_val_device, L_nz * sizeof(float));
    cudaMalloc((void **)&L_col_device, L_nz * sizeof(int));

    //拷贝L矩阵的空间
    cudaMemcpy(L_row_device, L_row_host, (A_size + 1) * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(L_val_device, L_val_host, L_nz * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(L_col_device, L_col_host, L_nz * sizeof(int), cudaMemcpyHostToDevice);

    //创建D+U
    int* DU_row_device = 0;
    float* DU_val_device = 0;
    int* DU_col_device = 0;

    cudaMalloc((void **)&DU_row_device, (A_size + 1) * sizeof(int));
    cudaMalloc((void **)&DU_val_device, DU_nz * sizeof(float));
    cudaMalloc((void **)&DU_col_device, DU_nz * sizeof(int));

    //拷贝D矩阵的空间
    cudaMemcpy(DU_row_device, DU_row_host, (A_size + 1) * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(DU_val_device, DU_val_host, DU_nz * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(DU_col_device, DU_col_host, DU_nz * sizeof(int), cudaMemcpyHostToDevice);

    //创建x数组
    float* x_val_device = 0;

    cudaMalloc((void **)&x_val_device, A_size * sizeof(float));
    
    cudaMemcpy(x_val_device, x_val_host, A_size * sizeof(float), cudaMemcpyHostToDevice);

    //创建b
    float* b_val_device = 0;

    cudaMalloc((void **)&b_val_device, A_size * sizeof(float));

    cudaMemcpy(b_val_device, b_val_host, A_size * sizeof(float), cudaMemcpyHostToDevice);
    
    //创建一个临时空间来存矩阵右端的信息
    float* r_val_device = 0;

    cudaMalloc((void **)&r_val_device, A_size * sizeof(float));

    cudaMemcpy(r_val_device, b_val_host, A_size * sizeof(float), cudaMemcpyHostToDevice);
    
    //创建cusparseHandle_t句柄
    cusparseHandle_t handle;
    cusparseCreate(&handle);

    //声明矩阵格式
    cusparseMatDescr_t descr;
    cusparseCreateMatDescr(&descr);

    //设定为普通矩阵，从0开始索引
    cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL);
    cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO);

    float alpha = -1;
    float beta = 1;

    //调用cusparseDcsrmv，来计算r = b - Lx_{k}
    cusparseScsrmv(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, A_size, A_size, L_nz, &alpha, descr, 
        L_val_device, L_row_device, L_col_device, x_val_device, &beta, r_val_device);

    //看看右侧矩阵的内容
    float* r_val_host = 0;

    cudaMallocHost((void **)&r_val_host, A_size * sizeof(float));

    cudaMemcpy(r_val_host, r_val_device, A_size * sizeof(float), cudaMemcpyDeviceToHost);
    
    // for(i = 0; i < A_size; i++){
    //     printf("%f\n", r_val_host[i]);
    // }

    //然后使用D+U，作为左侧的参数参与计算
    cusparseSolveAnalysisInfo_t info;
    cusparseCreateSolveAnalysisInfo(&info);

    //先进行分析
    cusparseScsrsv_analysis(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, A_size, nz, descr, DU_val_device,
         DU_row_device, DU_col_device, info);

    alpha = 1;

    //然后进行计算
    cusparseScsrsv_solve(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, A_size, &alpha, descr, DU_val_device, 
        DU_row_device, DU_col_device, info, r_val_device, x_val_device);

    //这里查看传回来的东西
    //将数据传回
    cudaMemcpy(x_val_host, x_val_device, A_size * sizeof(float), cudaMemcpyHostToDevice);

    //打印x矩阵
    for(i = 0; i < A_size; i++){
        printf("%f\n", x_val_host[i]);
    }


    //结束高斯赛德尔迭代

    //析构设备端矩阵
    cudaFree(L_row_device);
    cudaFree(L_val_device);
    cudaFree(L_col_device);

    cudaFree(DU_row_device);
    cudaFree(DU_val_device);
    cudaFree(DU_col_device);

    cudaFree(x_val_device);
    cudaFree(r_val_device);
    cudaFree(b_val_device);
    

    //析构主机端的A矩阵
    cudaFreeHost(A_col_host);
    cudaFreeHost(A_val_host);
    cudaFreeHost(A_row_host);

    //析构主机端的U矩阵
    cudaFreeHost(DU_row_host);
    cudaFreeHost(DU_val_host);
    cudaFreeHost(DU_col_host);

    //析构主机端的L矩阵
    cudaFreeHost(L_row_host);
    cudaFreeHost(L_val_host);
    cudaFreeHost(L_col_host);

    //析构x
    cudaFreeHost(x_val_host);

    //析构b
    cudaFreeHost(b_val_host);

    //析构r
    cudaFreeHost(r_val_host);
}